from typing import List, Dict, Optional, Tuple
import numpy as np
from sklearn.linear_model import LinearRegression
from sklearn.metrics import mean_squared_error, r2_score
from sklearn.preprocessing import StandardScaler
import traceback

from app.schemas.evolution_predict_schema import (
    AlgorithmResponse, OutputParams, TaskStatus,
    EvolutionPredictInputParams, EvolutionPredictOutputParams,
    TimePredictionResult
)
from app.utils.logger import logger
from app.tools.redis_client import redis_client

class EvolutionPredictService:
    def __init__(self):
        self.supported_methods = ["linear", "exp", "lstm"]
    
    def prepare_sequence_data(
        self, 
        historical_data: List[List[List[int]]], 
        window_size: int
    ) -> Tuple[np.ndarray, np.ndarray, StandardScaler]:
        """准备序列数据用于预测，增加数据标准化"""
        X, y = [], []
        flat_data = np.array(historical_data).reshape(len(historical_data), -1)
        
        # 数据标准化
        scaler = StandardScaler()
        scaled_data = scaler.fit_transform(flat_data)
        
        for i in range(len(scaled_data) - window_size):
            X.append(scaled_data[i:i+window_size])
            y.append(scaled_data[i+window_size])
        
        return np.array(X), np.array(y), scaler
    
    def linear_predict(
        self, 
        historical_data: List[List[List[int]]], 
        prediction_steps: int,
        window_size: int
    ) -> Tuple[List[List[List[int]]], List[float]]:
        """改进的线性回归预测方法"""
        X, y, scaler = self.prepare_sequence_data(historical_data, window_size)
        
        # 重塑数据以适应线性回归
        X_reshaped = X.reshape(X.shape[0], -1)
        y_reshaped = y.reshape(y.shape[0], -1)
        
        # 训练模型
        model = LinearRegression()
        model.fit(X_reshaped, y_reshaped)
        
        predictions = []
        confidence_scores = []
        last_window = np.array(historical_data[-window_size:])
        last_window_scaled = scaler.transform(last_window.reshape(window_size, -1))
        
        for _ in range(prediction_steps):
            # 预测下一个时间步
            next_pred_scaled = model.predict(last_window_scaled.reshape(1, -1))
            
            # 反向转换预测结果
            next_pred = scaler.inverse_transform(next_pred_scaled)
            next_pred = next_pred.reshape(last_window.shape[1:])
            
            # 约束预测结果为非负整数
            next_pred = np.maximum(np.round(next_pred), 0).astype(int)
            
            # 计算置信度分数并处理特殊值
            try:
                r2 = r2_score(y_reshaped, model.predict(X_reshaped))
                mse = mean_squared_error(y_reshaped, model.predict(X_reshaped))
                confidence = np.clip(r2 * (1 - np.sqrt(mse) / np.mean(y_reshaped)), 0, 1)
                # 确保confidence是有效的浮点数
                if np.isnan(confidence) or np.isinf(confidence):
                    confidence = 0.0
            except Exception:
                confidence = 0.0
            
            predictions.append(next_pred.tolist())
            confidence_scores.append(float(confidence))
            
            # 更新滑动窗口
            last_window = np.roll(last_window, -1, axis=0)
            last_window[-1] = next_pred
            last_window_scaled = scaler.transform(last_window.reshape(window_size, -1))
        
        return predictions, confidence_scores
    
    async def update_task_status(
        self, 
        task_id: str, 
        task_data: dict, 
        progress: int, 
        logs: str
    ):
        """更新任务状态"""
        try:
            response = {
                "task_id": task_id,
                "task_status": TaskStatus.RUNNING,
                "task_progress": progress,
                "task_logs": logs,
                "output_params": {  # 添加空的output_params
                    "evolution_predict_results": None
                }
            }
            await redis_client.update_data(f'{task_id}_result', response)
        except Exception as e:
            logger.error(f"更新任务状态失败: {str(e)}")
            raise
    
    async def run_evolution_predict(
        self, 
        input_params: dict, 
        task_id: str
    ) -> AlgorithmResponse:
        """运行演化预测算法的主函数"""
        task_data = await redis_client.get_data(task_id)
        
        try:
            # 验证输入参数
            if not input_params or "evolution_predict_params" not in input_params:
                raise ValueError("缺少必要的预测参数")
            
            predict_params = EvolutionPredictInputParams(**input_params.get("evolution_predict_params", {}))
            
            # 验证历史数据
            if not predict_params.historical_data or len(predict_params.historical_data) < predict_params.window_size:
                raise ValueError("历史数据不足或窗口大小设置不合理")
            
            await self.update_task_status(task_id, task_data, 10, "开始运行演化预测算法...")
            
            # 执行预测
            try:
                if predict_params.prediction_method == "linear":
                    predictions, confidence_scores = self.linear_predict(
                        predict_params.historical_data,
                        predict_params.prediction_steps,
                        predict_params.window_size
                    )
                    await self.update_task_status(task_id, task_data, 50, "预测计算完成...")
                else:
                    raise ValueError(f"不支持的预测方法: {predict_params.prediction_method}")
            except Exception as e:
                await self.update_task_status(
                    task_id, 
                    task_data, 
                    0, 
                    f"预测计算失败: {str(e)}"
                )
                raise
            
            # 构建预测结果
            try:
                time_results = []
                for t, (pred, conf) in enumerate(zip(predictions, confidence_scores)):
                    # 计算 MSE 时处理特殊值
                    try:
                        mse = float(mean_squared_error(
                            np.array(predict_params.historical_data[-1]), 
                            np.array(pred)
                        ))
                        if np.isnan(mse) or np.isinf(mse):
                            mse = float('inf')  # 使用一个合理的默认值
                    except Exception:
                        mse = float('inf')
                    
                    result = TimePredictionResult(
                        time_step=t + 1,
                        predicted_edges=pred,
                        confidence_score=float(conf),  # 确保是普通浮点数
                        prediction_metrics={"mse": mse}
                    )
                    time_results.append(result)
                
                await self.update_task_status(task_id, task_data, 80, "结果整理完成...")
                
                # 计算整体准确率
                overall_accuracy = np.mean(confidence_scores)
                
                evolution_output = EvolutionPredictOutputParams(
                    predictions=time_results,
                    parameters={
                        "prediction_steps": predict_params.prediction_steps,
                        "window_size": predict_params.window_size,
                        "method": predict_params.prediction_method
                    },
                    overall_accuracy=overall_accuracy
                )
                
                # 构建最终响应
                response = AlgorithmResponse(
                    task_id=task_id,
                    task_status=TaskStatus.COMPLETED,
                    task_progress=100,
                    task_logs="预测完成",
                    output_params=OutputParams(evolution_predict_results=evolution_output)
                )
                
                # 更新最终结果到Redis
                await redis_client.update_data(f'{task_id}_result', response.dict())
                
                return response
                
            except Exception as e:
                await self.update_task_status(
                    task_id, 
                    task_data, 
                    0, 
                    f"结果处理失败: {str(e)}"
                )
                raise
            
        except Exception as e:
            error_msg = f"演化预测算法执行失败: {str(e)}\n{traceback.format_exc()}"
            logger.error(error_msg)
            
            # 更新失败状态到Redis
            error_response = {
                "task_id": task_id,
                "task_status": TaskStatus.FAILED,
                "task_progress": 0,
                "task_logs": error_msg,
                "error_message": str(e),
                "output_params": {
                    "evolution_predict_results": None
                }
            }
            await redis_client.update_data(f'{task_id}_result', error_response)
            raise 